# -- coding: utf-8 --
"""
速度修正方程求解
"""

# third party lib
import numpy as np
# private lib
from utils.constant import InterpMode, DifferenceMode
from utils.numeric import grid_interp, difference


def correct_V(params: dict, p_prime_node: np.ndarray) -> np.ndarray:
    """ 速度修正计算
    Args:
        params: dict 计算需要的输入参
        p_prime_node: np.ndarray, 压力误差修正量
    Returns:
        V_g_apo: np.ndarray, 气相速度
        V_L_apo: np.ndarray, 液相速度
    """
    dt = params["dt"]
    dx = params["dx"]
    alpha_g_node = params["alpha_g_node"]
    alpha_g_segment = grid_interp(alpha_g_node, InterpMode.MID)
    alpha_L_node = params["alpha_L_node"]
    alpha_L_segment = grid_interp(alpha_L_node, InterpMode.MID)
    alpha_D_node = params["alpha_D_node"]
    alpha_D_segment = grid_interp(alpha_D_node, InterpMode.MID)
    
    rho_g_node = params["rho_g_node"]
    rho_g_segment = grid_interp(rho_g_node, InterpMode.MID)
    rho_L_node = params["rho_L_node"]
    rho_L_segment = grid_interp(rho_L_node, InterpMode.MID)
    
    partial_p_prime_segment = difference(p_prime_node, dx, DifferenceMode.CDMs)  # 偏 p_prime, 中心差分
    V_g_prime_segment = -partial_p_prime_segment * (alpha_g_segment + alpha_D_segment)/ (alpha_g_segment * rho_g_segment / dt)
    V_L_prime_segment = -partial_p_prime_segment* alpha_L_segment/ (rho_L_segment * alpha_L_segment / dt)
    
    V_g_segment = params["V_g_segment"]
    V_L_segment = params["V_L_segment"]
    V_g_segment += V_g_prime_segment
    V_L_segment += V_L_prime_segment
    return V_g_segment, V_L_segment
